/**
 * @file main.cpp
 *
 * Copyright (C) 2023-2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include <cmath>
#include <cstdint>
#include <iostream>
#include <thread>
#include <mpi.h>
#include "acl/acl.h"
#include "common.h"
#include "op_runner.h"

#include <sys/mman.h>
#include <asm-generic/mman-common.h>
#include <cstring>
#include "mem_share.h"

bool g_isDevice = false;
int deviceId = 0;
int RANK_DIM = std::atoi(std::getenv("RANK_SIZE"));
int used_core_num = std::atoi(std::getenv("USED_CORE_NUM"));
int p_value = std::atoi(std::getenv("P_VALUE"));
int p_value_comm = std::atoi(std::getenv("P_VALUE_COMM"));
int prompt_len = std::atoi(std::getenv("PROMPT_LEN"));
int buff_num = 8;
int call_times = std::atoi(std::getenv("CALL_TIMES"));
int SAVE_ONCE = std::atoi(std::getenv("IF_CHECK"));

OperatorDesc CreateOpDesc()
{
    // define operator
    int64_t M = prompt_len;
    int64_t K = 1024;
    int64_t N = 4096;
    int64_t rank_size = RANK_DIM;
    int64_t M_per_rank = M / rank_size;
    
    // used_core_num_ = 20;
    // loop_per_rank = p_value * used_core_num_ / rank_size = 5;
    // 5 * 128 * 256 * 8 * 2
    
    int64_t buffer_len = 128 * rank_size * ((p_value * used_core_num / rank_size + p_value_comm - 1) / p_value_comm) * 2;

    std::vector<int64_t> shapeA{M, K};
    std::vector<int64_t> shapeB{K, N};
    std::vector<int64_t> shapeScale{N};
    std::vector<int64_t> shapePertokenScale{M};
    std::vector<int64_t> shapeC{M_per_rank, N};

    std::vector<int64_t> shapeBuffer{buffer_len + 1,256 * p_value_comm};

    aclDataType dataTypeA = ACL_INT8;
    aclDataType dataTypeB = ACL_INT8;
    aclDataType dataTypeScale = ACL_BF16;
    aclDataType dataTypeBias = ACL_BF16;
    aclDataType dataTypePertokenScale = ACL_FLOAT;
    aclDataType dataTypeC = ACL_BF16;
    aclFormat format_nd = ACL_FORMAT_ND;
    aclFormat format_nz = ACL_FORMAT_FRACTAL_NZ;

    OperatorDesc opDesc;

    opDesc.AddInputTensorDesc(dataTypeA, shapeA.size(), shapeA.data(), format_nd);
    opDesc.AddInputTensorDesc(dataTypeB, shapeB.size(), shapeB.data(), format_nz);
    opDesc.AddInputTensorDesc(dataTypeScale, shapeScale.size(), shapeScale.data(), format_nd);
    opDesc.AddInputTensorDesc(ACL_DT_UNDEFINED, 0, nullptr, ACL_FORMAT_UNDEFINED);
    opDesc.AddInputTensorDesc(dataTypePertokenScale, shapePertokenScale.size(), shapePertokenScale.data(), format_nd);
    opDesc.AddOutputTensorDesc(dataTypeC, shapeC.size(), shapeC.data(), format_nd);

    for(int rank = 0; rank < buff_num; rank++) {
    opDesc.AddOutputTensorDesc(dataTypeC, shapeBuffer.size(), shapeBuffer.data(), format_nd);
    }

    // 构造算子属性
    auto opAttr = opDesc.opAttr;
    aclopSetAttrInt(opAttr, "output_type", dataTypeC);
    aclopSetAttrBool(opAttr, "transpose_a", false);
    aclopSetAttrBool(opAttr, "transpose_b", false);
    aclopSetAttrInt(opAttr, "split_trans_num", 0);
    aclopSetAttrInt(opAttr, "rank_size", rank_size);
    aclopSetAttrInt(opAttr, "rank", 0);
    aclopSetAttrInt(opAttr, "p_value", 2);

    return opDesc;
}

bool SetInputData(OpRunner &runner)
{
    size_t fileSize = 0;
    
    ReadFile("../input/input_a.bin", fileSize, runner.GetInputBuffer<void>(0), runner.GetInputSize(0));
    ReadFile("../input/input_b.bin", fileSize, runner.GetInputBuffer<void>(1), runner.GetInputSize(1));
    //ReadFile("../input/input_bias.bin", fileSize, runner.GetInputBuffer<void>(2), runner.GetInputSize(2));
    ReadFile("../input/input_scale.bin", fileSize, runner.GetInputBuffer<void>(2), runner.GetInputSize(2));
    ReadFile("../input/input_pertoken_scale.bin", fileSize, runner.GetInputBuffer<void>(4), runner.GetInputSize(4));

    // ReadFile("../input/buff_tmp_0.bin", fileSize, runner.GetOutputBuffer<void>(1), runner.GetOutputSize(1));
    // ReadFile("../input/buff_tmp_1.bin", fileSize, runner.GetOutputBuffer<void>(2), runner.GetOutputSize(2));
    // ReadFile("../input/buff_tmp_2.bin", fileSize, runner.GetOutputBuffer<void>(3), runner.GetOutputSize(3));
    // ReadFile("../input/buff_tmp_3.bin", fileSize, runner.GetOutputBuffer<void>(4), runner.GetOutputSize(4));
    // ReadFile("../input/buff_tmp_4.bin", fileSize, runner.GetOutputBuffer<void>(5), runner.GetOutputSize(5));
    // ReadFile("../input/buff_tmp_5.bin", fileSize, runner.GetOutputBuffer<void>(6), runner.GetOutputSize(6));
    // ReadFile("../input/buff_tmp_6.bin", fileSize, runner.GetOutputBuffer<void>(7), runner.GetOutputSize(7));
    // ReadFile("../input/buff_tmp_7.bin", fileSize, runner.GetOutputBuffer<void>(8), runner.GetOutputSize(8));


    INFO_LOG("Set input success");
    return true;
}

bool ProcessOutputData(OpRunner &runner, uint32_t rankId)
{
    WriteFile("../output/output_" + std::to_string(rankId) + ".bin", runner.GetOutputBuffer<void>(0), runner.GetOutputSize(0));
    // WriteFile("../output/output_" + std::to_string(rankId) + ".bin", runner.GetOutputBufferDevOutput0<void>(), runner.GetOutputSize(0));
    WriteFile_bf16Toft32("../output/output_" + std::to_string(rankId) + "_ft32.bin", runner.GetOutputBuffer<void>(0), runner.GetOutputSize(0));
    INFO_LOG("Write output success");

    return true;
}


bool InitResource()
{
    std::string output = "../output";
    if (access(output.c_str(), 0) == -1) {
        int ret = mkdir(output.c_str(), 0700);
        if (ret == 0) {
            INFO_LOG("Make output directory successfully");
        } else {
            ERROR_LOG("Make output directory fail");
            return false;
        }
    }

    if (aclInit(nullptr) != ACL_SUCCESS) {
        ERROR_LOG("acl init failed");
        return false;
    }

    // for (int32_t i = 0; i < RANK_DIM; i++) {
    //     if (aclrtSetDevice(i) != ACL_SUCCESS) {
    //         ERROR_LOG("Set device failed. deviceId is %u", i);
    //         for (uint32_t j = 0; j < i; j++) {
    //             (void)aclrtResetDevice(j);
    //         }
    //         (void)aclFinalize();
    //         return false;
    //     }
    // }
    return true;
}

bool RunOp(uint32_t rankId)  
{

    // create contest
    if (aclInit(nullptr) != ACL_SUCCESS) {
        ERROR_LOG("acl init failed");
        return false;
    }

    if (aclrtSetDevice(rankId) != ACL_SUCCESS) {
        ERROR_LOG("Set device failed. deviceId is %u", rankId);
        (void)aclrtResetDevice(rankId);
        (void)aclFinalize();
        return false;
    }

    aclrtContext context;
    if (aclrtCreateContext(&context, rankId) != ACL_SUCCESS) {
        ERROR_LOG("Create context failed. deviceId is %u", rankId);
        (void)aclrtResetDevice(rankId);
        (void)aclFinalize();
        return false;
    }
    // create stream
    aclrtStream stream;
    if (aclrtCreateStream(&stream) != ACL_SUCCESS) {
        ERROR_LOG("Create stream failed. deviceId is %u", rankId);
        (void)aclrtDestroyContext(context);
        (void)aclrtResetDevice(rankId);
        (void)aclFinalize();
        return false;
    }

    // set context
    if (aclrtSetCurrentContext(context) != ACL_SUCCESS) {
        ERROR_LOG("Set current context failed, deviceId=%u", rankId);
        (void)aclrtDestroyStream(stream);
        (void)aclrtDestroyContext(context);
        (void)aclrtResetDevice(rankId);
        (void)aclFinalize();
        return false;
    }

    // runMode is ACL_HOST which represents app is running in host
    // runMode is ACL_DEVICE which represents app is running in device
    aclrtRunMode runMode;
    if (aclrtGetRunMode(&runMode) != ACL_SUCCESS) {
        ERROR_LOG("Get run mode failed");
        (void)aclrtDestroyStream(stream);
        (void)aclrtDestroyContext(context);
        (void)aclrtResetDevice(rankId);
        (void)aclFinalize();
        return false;
    }

    std::cout <<"g_isDevice:" << g_isDevice << std::endl;
    g_isDevice = (runMode == ACL_DEVICE);
    INFO_LOG("Get RunMode[%d] success", runMode);
    if(RANK_DIM>1){
    	MPI_Barrier(MPI_COMM_WORLD);
    }

    // create op desc
    OperatorDesc opDesc = CreateOpDesc(); 

    // create Runner
    OpRunner opRunner(&opDesc, rankId);
    if (!opRunner.Init()) {
        ERROR_LOG("Init OpRunner failed");
        (void)aclrtDestroyStream(stream);
        (void)aclrtDestroyContext(context);
        (void)aclrtResetDevice(rankId);
        (void)aclFinalize();
        return false;
    }

    // Load inputs
    if (!SetInputData(opRunner)) {  //除非不同卡读不同的数据，否则不用修改
        ERROR_LOG("Set input data failed");
        (void)aclrtDestroyStream(stream);
        (void)aclrtDestroyContext(context);
        (void)aclrtResetDevice(rankId);
        (void)aclFinalize();
        return false;
    }

    // Run op
    for(int j=0; j<call_times; j++){
        if (!opRunner.RunOp(stream)) { 
            ERROR_LOG("Run op failed");
            (void)aclrtDestroyStream(stream);
            (void)aclrtDestroyContext(context);
            (void)aclrtResetDevice(rankId);
            (void)aclFinalize();
            return false;
        }
        
        if (!SAVE_ONCE) {
            continue;
        }
        opRunner.SaveOutput(j);
        SAVE_ONCE = 0;
    }

    // process output data
    // if (!ProcessOutputData(opRunner, rankId)) { 
    //     ERROR_LOG("Process output data failed");
    //     (void)aclrtDestroyStream(stream);
    //     (void)aclrtDestroyContext(context);
    //     (void)aclrtResetDevice(rankId);
    //     (void)aclFinalize();
    //     return false;
    // }

    INFO_LOG("Run op success");
    
    (void)aclrtDestroyStream(stream);
    (void)aclrtDestroyContext(context);
    (void)aclrtResetDevice(rankId);
    (void)aclFinalize();
    return true;
}

int main(int argc, char **argv)
{
    if(RANK_DIM > 1){
        MPI_Init(NULL, NULL);
        int rank, world_size;
        MPI_Comm_size(MPI_COMM_WORLD, &world_size);
        MPI_Comm_rank(MPI_COMM_WORLD, &rank);
        shared_mm_init(world_size);
        MPI_Barrier(MPI_COMM_WORLD);

        RunOp(rank);
        MPI_Barrier(MPI_COMM_WORLD);
        MPI_Finalize();
    }else {
        RunOp(0);
    }


    return SUCCESS;

}
